import matplotlib
import torch
import torchvision
import torch.utils
from torch import nn
from d2l import torch as d2l

# 图像增广
# d2l.set_figsize()
# img = d2l.Image.open('D:\\PytorchLearn\\pytorch_learning\\Project1\\capture_20221115110104367.bmp')
# d2l.plt.imshow(img)
#
# def apply(img,aug,num_rows=2,num_cols=4,scale=1.5):
#     Y = [aug(img) for _ in range(num_rows *num_cols)]
#     d2l.show_images(Y,num_rows,num_cols,scale)
#
# apply(img,torchvision.transforms.RandomHorizontalFlip()) # 左右反转图像
# apply(img,torchvision.transforms.RandomVerticalFlip()) # 上下反转
# shape_aug = torchvision.transforms.RandomResizedCrop(
#     (200,100),scale=(0.1,1),ratio=(0.5,2))
# apply(img,shape_aug) # 随即裁剪
# apply(img,torchvision.transforms.ColorJitter(
#     brightness=0.5,contrast=0,saturation=0,hue=0)) # 随机更改亮度
# apply(img,torchvision.transforms.ColorJitter(
#     brightness=0,contrast=0,saturation=0,hue=0.5)) # 随机改变色调
# apply(img,torchvision.transforms.ColorJitter(
#     brightness=0.5,contrast=0.5,saturation=0.5,hue=0.5))  # 亮度 对比度，饱和度，色调

# all_images = torchvision.datasets.CIFAR10(train=True,root='./dataset',download=True)
# d2l.show_images([
#     all_images[i][0] for i in range(32)],4,8,scale=0.8)

train_augs = torchvision.transforms.Compose([
    torchvision.transforms.RandomHorizontalFlip(),
    torchvision.transforms.ToTensor()])
test_augs = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

def load_cifar10(is_train,augs,batch_size):
    dataset =torchvision.datasets.CIFAR10(train=is_train,root='./dataset',transform=augs,download=True)
    dataloader = torch.utils.data.DataLoader(
        dataset,batch_size=batch_size,shuffle=is_train,
        num_workers = 2)
    return dataloader

def train_batch_ch13(net,X,y,loss,trainer):
    net.train()
    trainer.zero_grad()
    pred=net(X)
    l = loss(pred,y)
    l.sum().backward
    trainer.step()
    train_loss_sum = l.sum()
    train_acc_sum = d2l.accuracy(pred,y)
    return train_loss_sum,train_acc_sum







